baby-llama2-chinese 实践笔记

baby-llama2-chinese 项目，演示了如何从零训练一个参数量 500M-1B 的 LLM。该项目完整包含了训练、SFT指令微调、奖励模型等过程，是一个很好的教学项目。

环境搭建：首先从 GitHub 上 clone 项目。相关依赖包可参见 requirements.txt。

在本文中，记录了我对该项目的自学过程。该过程同时参考了《jjchen1》的高质量笔记。

Note

文中提到一些预训练语料，我统一收集至笔记《大模型数据集收集》中。

下载运行预训练微调模型

我先从预训练模型下手。作者同时提供了 base 和微调模型，我选择 Llama2-Chinese-218M-v1-NormalChat。将下载好的模型放置到项目的 out 目录下。

修改 eval.py 中的 ckpt_path：

ckpt_path = "out/baby-llama2-chinese/Llama2-Chinese-218M-v2/sft_model_norm_epoch5.pth"

注意参数要对上：

max_seq_len=1024
dim=1024
n_layers=12
n_heads=8

运行 eval.py：

(.venv) [maxiee@archlinux baby-llama2-chinese]$ python eval.py 
[prompt]: 最近我在办公室坐久了会感到头晕，请问这是什么原因?有什么缓解办法吗？
[answer]:  我无法确定您所描述的问题，因此无法回答。是一种完成的工作或任务，您需要提供更多信息或采取任何可能蛋白质时间或建议来帮助您处理。如果您需要我提供信息，请告诉我您需要答案指出了哪些原因。
---------------
[prompt]: 前列腺囊肿的症状是什么？
[answer]:  您是否对您所推荐购买的决策具体建议？如果可以的话，请提供更多信息。
---------------
[prompt]: 请问，世界上最大的动物是什么？
[answer]:  世界上最大的动物是蓝鲸。蓝鲸段头鲸，也称为蓝鲸河见鲸，是雨海、海洋中的垃圾袋鲸，是大型哺乳动物，身长约为40米，体重超过30000公斤。

2 jjchen1 笔记

我是从 jjchen1 大佬的笔记中得知 baby-llama2-chinese 项目的。jjchen1 的笔记是一个为期 30 天的 LLM 的探索，这里摘录与本文项目相关的文章，其他文章同样值得推荐。

Day 6 - Baby LLama2：

读完笔记后，我相当地佩服 jjchen1 大佬！

4 开始预训练！

接下来，该我出场了，在大佬们的指点下，我也准备走一遍这个流程。

baby-llama2-chinese 作者提供了加工好的语料，共计634亿Tokens的数据量，文件大小为 118.18GB。我直接复用这个语料。

为了加快训练速度，减小模型参数：

相较于 92M 模型，层数由8减为4，多头注意力由8减为4。

修改了一下语料加载代码，一次性加载目录下所有 .bin：

data_root_path = "/mnt/ssd2/home/maxiee/LLM/baby-llama2-chinese/pretrain_data_bin"
# 搜索 data_root_path 下所有 .bin 文件，保存至 data_path_list
data_path_list = []
for root, dirs, files in os.walk(data_root_path):
    for file in files:
        if file.endswith(".bin"):
            print(os.path.join(root, file))
            data_path_list.append(os.path.join(root, file))

jjchen1 笔记中有提到：memmap=False，防止内存干爆，我实测了以下：

train_ds = PretrainDataset(data_path_list, max_length=max_seq_len, memmap=False)

执行预训练，我是单卡：

torchrun --standalone pretrain.py

效果赏析（挑选一些看起来好一点的）：

Enter a question: 信念不是到处去
[prompt]: 信念不是到处去
[answer]: 打工，反而能像我一样学习

[prompt]: 什么叫作失败?失败是
[answer]: 必须的。
你为什么认为胜利呢？我为什么要提出胜利或失败，你为什么觉得败，你为什么要失败。成功不是为了失败和失败，但当一个人遇到成功时，那人可能不是他自己的命运，而是当一个人去。结果就失败了。结果就是失败了。他不能坚持失败。于是又把胜利归告了。这样，失败就了。结果就败了。。

[prompt]: 得到智慧的惟一办法，就是用
[answer]: 我们的智慧和智慧去智慧。。主要产品：《智障》系列图书简介：【系列构成】
智障系列由北京智障科技股份有限公司、深圳智障科技发展有限公司倾力打造，是中国智障科技旗下专业研究人员，是智障科技旗下公司运营的专业研究团队。也是公司所有咨询专家和顾问。智障人从咨询、咨询、咨询到技术咨询；到品牌规划：提供、指导、执行，帮助客户实现营销策略。。编辑推荐：

[prompt]: 我从来不怕得罪朋友，因为
[answer]: 我是大人物。

预训练代码阅读

训练大模型的第一步是数据处理（data_process.py），第二步是正式预训练（pretrain.py)。本节中阅读正式预训练（pretrain.py)部分的代码。

首先，有众多超参数：

参数 说明 理解
out_dir 输出目录，用于保存训练过程中生成的模型检查点和日志。
max_epoch 最大训练轮次。每个轮次都会遍历一次整个训练数据集。
eval_interval 评估间隔，每训练eval_interval轮后，模型会在验证集上进行评估。
log_interval 日志间隔，每训练log_interval轮后，训练的进度和性能指标会被记录下来。
save_interval 保存间隔，每训练save_interval轮后，模型的当前状态会被保存下来。
eval_iters 评估迭代次数，定义了在每次评估过程中，模型需要处理的数据批次数量。
eval_only 如果为True，脚本在第一次评估后就会退出。
always_save_checkpoint 如果为True，每次评估后都会保存一个模型检查点。
init_from 模型初始化方式，可以从头开始（'scratch'），从上次的检查点恢复（'resume'），或者从预训练的GPT2模型开始（'gpt2*'）。
gradient_accumulation_steps 梯度累积步数，用于模拟更大的批次大小。
batch_size 批次大小，如果gradient_accumulation_steps大于1，这就是微批次的大小。
max_seq_len 序列的最大长度，即模型可以处理的输入序列的最大长度。
dim 模型的维度，即模型内部表示的大小。
n_layers Transformer模型的层数。
n_heads Transformer模型的注意力头的数量。
multiple_of 用于确保序列长度是某个数的倍数，这对某些硬件加速器（如GPU）的性能优化很重要。
dropout Dropout率，用于防止模型过拟合。在预训练阶段，通常设置为0，在微调阶段，可以尝试设置为0.1或更高。
bias 是否在LayerNorm和Linear层中使用偏置。
learning_rate 学习率，决定了模型参数在每次更新时的调整幅度。如果学习率过大，训练可能会不稳定；如果学习率过小，训练可能会过慢。
weight_decay 权重衰减，用于防止模型过拟合。它通过在优化器的更新规则中添加一个与当前参数值成比例的项来实现。
beta1, beta2 这是Adam优化器的超参数，用于计算梯度和梯度平方的移动平均值。
grad_clip 梯度裁剪值，用于防止梯度爆炸。如果梯度的范数超过这个值，就会被裁剪。
decay_lr
warmup_iters
lr_decay_iters
min_lr		 这些变量用于配置学习率衰减策略。decay_lr决定是否使用学习率衰减，warmup_iters是预热步数，在这些步数内，学习率会线性增加到learning_rate，lr_decay_iters是学习率开始衰减的步数，min_lr是学习率的最小值。
backend 分布式数据并行（DDP）的后端，可以是'nccl'，'gloo'等。
device 训练设备，可以是'cpu'，'cuda'，'cuda:0'，'cuda:1'等。
dtype dtype: 数据类型，可以是'float32'，'bfloat16'，或'float16'。 如果选择'float16'，将自动实现梯度缩放，这对于防止浮点数精度问题很有用。
compile 是否使用PyTorch 2.0编译模型以提高速度。

该脚本首先判断是否运行于分布式环境下（由 ddp 变量标识），这里我主要看单 GPU 的场景。

接下来计算每次迭代处理的 token 数量：

tokens_per_iter = 
	# 梯度累积步数。
	gradient_accumulation_steps \
	# 分布式数据并行世界大小
	* ddp_world_size \
	# 批次大小
	* batch_size \
	# 序列的最大长度
	* max_seq_len

其中：

在我的场景下，就是 max_seq_len * batch_size，也就是一次喂给 GPU 的 Token 数量，即迭代（iterator）。

继续进行设置：

# 设置了 PyTorch 的随机数生成器的种子。
torch.manual_seed(1337 + seed_offset)

# 启用了 TensorFlow 32位浮点数（TF32）格式。
# TF32 是 NVIDIA 在其 Ampere 架构的 GPU 中引入的一种新的浮点数格式
# 它可以在不牺牲太多精度的情况下提高计算速度。
torch.backends.cuda.matmul.allow_tf32 = True  # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True  # allow tf32 on cudnn

# 设置了设备类型为 "cuda"，表示将在 GPU 上进行计算。
device_type = "cuda"

# 根据 `dtype` 的值选择了一个 PyTorch 的数据类型。
# 如果 `dtype` 是 "float32"，`ptdtype` 就会被设置为 `torch.float32`；
# 如果 `dtype` 是 "bfloat16"，`ptdtype` 就会被设置为 `torch.bfloat16`；
# 如果 `dtype` 是 "float16"，`ptdtype` 就会被设置为 `torch.float16`。
# note: float16 data type will automatically use a GradScaler
ptdtype = {"float32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[dtype]

# 如果 `device_type` 是 "cuda"，`ctx` 就会被设置为 `torch.cuda.amp.autocast()`
# 这是一个自动混合精度（Automatic Mixed Precision，简称 AMP）的上下文管理器
# 它可以自动选择合适的数据类型进行计算，以提高计算速度和减少内存消耗。
ctx = (
    nullcontext()
    if device_type == "cpu"
    else torch.cuda.amp.autocast()
)
# 用于存储最佳的验证损失
# 初始值被设置为一个非常大的数
# 这样在训练过程中只要遇到一个比这个数小的验证损失，就会更新
best_val_loss = 1e9

下面是加载处理后的 Token 语料，我对源代码做了修改，能够自动遍历目录下所有语料文件（.bin）：

data_root_path = "/mnt/ssd2/home/maxiee/LLM/baby-llama2-chinese/pretrain_data_bin"
# 搜索 data_root_path 下所有 .bin 文件，保存至 data_path_list
data_path_list = []
for root, dirs, files in os.walk(data_root_path):
    for file in files:
        if file.endswith(".bin"):
            print(os.path.join(root, file))
            data_path_list.append(os.path.join(root, file))

接下来创建预训练数据集类，以及训练采样器：

# 自定义的数据集类，它从 `data_path_list` 中的文件中加载数据
# 每个数据样本的最大长度为 `max_seq_len`。
# 如果 `memmap` 参数为 `True`，则使用内存映射文件来加载数据，这可以减少内存消耗。
train_ds = PretrainDataset(data_path_list, max_length=max_seq_len, memmap=True)

# DistributedSampler 是 PyTorch 的一个工具，用于在分布式训练中对数据进行采样。
train_sampler = torch.utils.data.distributed.DistributedSampler(train_ds)

# DataLoader 是 PyTorch 的一个工具，用于在训练过程中加载数据。
# 它可以自动地将数据分成多个批次，并在需要时加载每个批次的数据。
# 这里的参数设置如下：
train_loader = torch.utils.data.DataLoader(
    train_ds,
    batch_size=batch_size,
    pin_memory=False,
    drop_last=False,
    shuffle=False,        
    num_workers=0 if os.name == 'nt' else 4,
    sampler=train_sampler
)

其中，在 PretrainDataset 部分，jjchen1 在 Day 13 - Baby LLama2 Chinese (7) 中有提及：原因是，data_path_list 是个列表，允许传入多个语料文件。但是，当采用 memmap=True 时，PretrainDataset 只会用 np.memmap 加载第一个文件。这里，根据源作者提供的处理后语料能够看出，他是将所有原始语料集拼接为一个语料文件进行训练。

当 memmap=False 时，需要将所有语料 .bin 都加载到内存中，这是非常占用内存的。并且在语料集非常大时，这一步耗时也很长，需要等待一段时间才能开始训练。

shuffle 被设置为 False，这意味着数据不会在每个 epoch 开始时被随机打乱。

下面创建模型实例：

model=init_model()
model.to(device)

如下是 init_model 的具体实现。通过 init_from 变量判断模型的 3 种模式：scratch 从头训练、resume 增量预训练。

def init_model():
    # 模型参数
    model_args = dict(
        dim=dim,           # 维度
        n_layers=n_layers, # 层数
        n_heads=n_heads,   # 头数
        n_kv_heads=n_heads,# 头数
        vocab_size=64793,  # 词汇表大小
        multiple_of=multiple_of,
        max_seq_len=max_seq_len, # 序列的最大长度
        dropout=dropout,
    )  # start with model_args from command line
    if init_from == "scratch":
	    # 从头训练
        # init a new model from scratch
        print("Initializing a new model from scratch")
        gptconf = ModelArgs(**model_args)
        model = Transformer(gptconf)
    elif init_from == "resume":
	    # 增量预训练
        print(f"Resuming training from {out_dir}")
        # resume training from a checkpoint.
        ckpt_path = os.path.join(out_dir, "ckpt.pt")
        # 加载检查点文件
        checkpoint = torch.load(ckpt_path, map_location=device)
        checkpoint_model_args = checkpoint["model_args"]
        # force these config attributes to be equal otherwise we can't even resume training
        # the rest of the attributes (e.g. dropout) can stay as desired from command line
        for k in ["dim", "n_layers", "n_heads", "n_kv_heads", "vocab_size", "multiple_of", "max_seq_len"]:
            model_args[k] = checkpoint_model_args[k]
        # create the model
        gptconf = ModelArgs(**model_args)
        model = Transformer(gptconf)
        state_dict = checkpoint["model"]
        # fix the keys of the state dictionary :(
        # honestly no idea how checkpoints sometimes get this prefix, have to debug more
        unwanted_prefix = "_orig_mod."
        for k, v in list(state_dict.items()):
            if k.startswith(unwanted_prefix):
                state_dict[k[len(unwanted_prefix) :]] = state_dict.pop(k)
        model.load_state_dict(state_dict)
        iter_num = checkpoint["iter_num"]
        best_val_loss = checkpoint["best_val_loss"]
    return model

接下来创建优化器并开始训练：

# initialize a GradScaler. If enabled=False scaler is a no-op
scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
# optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)

# ...

raw_model = model.module if ddp else model # unwrap DDP container if needed
# training loop

# 计算了每个训练周期（epoch）中的迭代次数（iteration）
# 返回的是训练数据加载器 `train_loader` 中的批次数量。
# 在 PyTorch 中，`DataLoader` 对象是一个可迭代的对象，它将数据集分成多个批次，
# 每个批次包含多个数据样本。
# 当我们对 `DataLoader` 对象使用 `len()` 函数时，返回的是批次的数量
# 也就是每个训练周期需要进行的迭代次数。
iter_per_epoch=len(train_loader)
for epoch in range(max_epoch):
    train_epoch(epoch)

# 保存 epoch 训练结果            torch.save(raw_model.state_dict(),'{}/epoch_{}.pth'.format(save_dir,epoch))
    else:
        torch.save(raw_model.state_dict(),'{}/epoch_{}.pth'.format(save_dir,epoch))

看一个周期内是如何训练的 train_epoch：

def train_epoch(epoch):
    start_time=time.time()
    # 取出一个批次，根据 X 预测 Y
    for step, (X, Y) in enumerate(train_loader):
	    # 数据放到 GPU
        X=X.to(device)
        Y=Y.to(device)
        lr = get_lr(epoch*iter_per_epoch+step) if decay_lr else learning_rate
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
        # and using the GradScaler if data type is float16
        # for micro_step in range(gradient_accumulation_steps):
        # ……
        with ctx:
		    # 生成预测
            logits = model(X, Y)
            loss = raw_model.last_loss
            loss = loss / gradient_accumulation_steps
        # immediately async prefetch next batch while model is doing the forward pass on the GPU
        # backward pass, with gradient scaling if training in fp16
        scaler.scale(loss).backward()
        #
        if (step + 1) % gradient_accumulation_steps == 0:
            # clip the gradient
            if grad_clip != 0.0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
            # step the optimizer and scaler if training in fp16
            scaler.step(optimizer)
            scaler.update()
            # flush the gradients as soon as we can, no need for this memory anymore
            optimizer.zero_grad(set_to_none=True)
        # 打印日志
        if step % log_interval == 0:
            # 打印训练过程日志
        #
        if step % save_interval == 0:
	        # 保存中间过程
	        # 将模型切换到评估模式（evaluation mode）。
	        # 在评估模式下，某些类型的层（如 Dropout 和 BatchNorm）会改变它们的行为。
	        # 例如，Dropout 层在训练模式下会随机丢弃一部分神经元
	        # 而在评估模式下则不会丢弃任何神经元。
	        model.eval()
	        # 保存模型
		    torch.save(
			    model.module.state_dict(),
			    '{}/iter_{}.pth'.format(
				    save_dir,
				    int(step+epoch*iter_per_epoch)))
			# 这行代码将模型切换回训练模式。
			model.train()

6 Transformer 模型代码阅读

在本节中，阅读项目中使用到的 Transformer 模型代码。其中，我用 karpathy/minGPT 作为对照，来对比两者的差异。

TransformerBlock

一个 Transformer 块，Transformer 是由多层组成的（由层数 n_layers）控制，而每一层都是由 TransformerBlock 叠加而成。

TransformerBlock 块的实现，首先，输入数据 x 通过 self.attention_norm 进行归一化处理，然后传入 self.attention 层。注意力层的输出与原始输入 x 相加，得到 h。h 通过 self.ffn_norm 进行归一化处理，然后传入 self.feed_forward 层。前馈神经网络层的输出与 h 相加，得到最终的输出 out。具体代码如下：

class TransformerBlock(nn.Module):
    def __init__(self, layer_id: int, args: ModelArgs):
        super().__init__()
        self.n_heads = args.n_heads
        self.dim = args.dim
        self.head_dim = args.dim // args.n_heads
        self.attention = Attention(args)
        self.feed_forward = FeedForward(
            dim=args.dim,
            hidden_dim=4 * args.dim,
            multiple_of=args.multiple_of,
            dropout=args.dropout,
        )
        self.layer_id = layer_id
        self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
        self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

    def forward(self, x, freqs_cos, freqs_sin):
        h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
        out = h + self.feed_forward.forward(self.ffn_norm(h))
        return out

对比与 minGPT 的代码，Block 实现基本是一致的。FeedForward 也基本一致，但是存在少许差异。

需要注意，FeedForward 层的隐藏状态维度是 dim 的 4 倍，

FeedForward 是一个前馈神经网络它包含三个线性层和一个 dropout 层。初始化过程：

class FeedForward(nn.Module):
	def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
	    super().__init__()
	    # 将 `hidden_dim` 的值变为原来的 2/3。这是为了在后续的线性层中减少参数的数量。
	    hidden_dim = int(2 * hidden_dim / 3)
	    # multiple_of 可以控制让隐藏状态变大
	    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
	    # 这是第一个线性层，它将输入的维度从 dim 变为 hidden_dim。
	    self.w1 = nn.Linear(dim, hidden_dim, bias=False)
	    # 这是第二个线性层，它将输入的维度从 `hidden_dim` 变回 `dim`。
	    self.w2 = nn.Linear(hidden_dim, dim, bias=False)
	    # 这是第三个线性层，它将输入的维度从 `dim` 变为 `hidden_dim`。
	    self.w3 = nn.Linear(dim, hidden_dim, bias=False)
	    self.dropout = nn.Dropout(dropout)

def forward(self, x):
	return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))

forward 步骤拆解：

  1. self.w1(x): 输入数据 x 首先通过一个线性层 self.w1。这个线性层会对 x 进行线性变换。
  2. F.silu(...): 然后，线性层的输出通过 Sigmoid Linear Unit (SiLU) 激活函数。SiLU 函数也被称为 Swish 函数，它的公式是 silu(x) = x * sigmoid(x)，其中 sigmoid(x) 是 logistic sigmoid 函数。SiLU 函数可以增加模型的非线性，使得模型能够学习更复杂的模式。
  3. ... * self.w3(x): 接着，SiLU 函数的输出与 self.w3(x) 的结果相乘。这是一种称为门控（gating）的技术，它可以控制信息的流动。
  4. self.w2(...): 然后，乘法的结果通过另一个线性层 self.w2。这个线性层会对数据进行另一次线性变换。
  5. self.dropout(...): 最后，线性层的输出通过一个 dropout 层。在训练过程中，dropout 层会随机丢弃一部分数据，这可以防止模型过拟合。

与 minGPT 对比，两者使用的激活函数不同，两者对比：

特性 本文 minGPT
隐藏层维度 动态计算 输入维度 * 4
激活函数 SiLU (Swish) NewGELU (GELU 变体)

SiLU (Swish) 和 NewGELU (GELU 变体) 的对比:

总结:

建议:

参考资料:

Attention

在 forward 方法中，定义了模型的前向传播过程。首先，输入数据 x 被分解为查询、键和值。然后，这些查询和键被应用了旋转位置嵌入（Rotary Position Embedding，RoPE）。接着，键和值被复制以匹配查询的数量。然后，查询、键和值被重塑以便进行注意力计算。

具体实现如下：

def forward(
    self,
    x: torch.Tensor,
    freqs_cos: torch.Tensor,
    freqs_sin: torch.Tensor,
):
    bsz, seqlen, _ = x.shape

    # QKV
    # 首先，输入数据 `x` 被分解为查询、键和值。
    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
    xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
    xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
    xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

    # RoPE relative positional embeddings
    # 然后，这些查询和键被应用了旋转位置嵌入（Rotary Position Embedding，RoPE）。
    xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)

    # grouped multiquery attention: expand out keys and values
    # ...
    
    # make heads into a batch dimension
    xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
    xk = xk.transpose(1, 2)
    xv = xv.transpose(1, 2)

    # flash implementation
    if self.flash:
        output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
    else:
        # manual implementation
        scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
        assert hasattr(self, 'mask')
        scores = scores + self.mask[:, :, :seqlen, :seqlen]   # (bs, n_local_heads, seqlen, cache_len + seqlen)
        scores = F.softmax(scores.float(), dim=-1).type_as(xq)
        scores = self.attn_dropout(scores)
        output = torch.matmul(scores, xv)  # (bs, n_local_heads, seqlen, head_dim)

    # restore time as batch dimension and concat heads
    output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

    # final projection into the residual stream
    output = self.wo(output)
    output = self.resid_dropout(output)
    return output

其中，在预训练时，n_kv_heads 我设置为与 n_heads 相等。

当self.flash为真时，代码使用PyTorch的torch.nn.functional.scaled_dot_product_attention函数来执行缩放点积注意力计算。这是一种更快、更简洁的实现方式，通常用于优化性能和计算效率。

当self.flash为假时，则采用手动实现的注意力机制，这涉及到更详细的步骤，包括计算查询（xq）和键（xk）之间的得分，应用掩码（mask），执行softmax操作，然后计算最终的输出。这种实现方式提供了更多的自定义性和控制性，但可能不如torch.nn.functional.scaled_dot_product_attention函数高效。

flash 的判断是根据 PyTorch 中是否有 'scaled_dot_product_attention' 函数来判断的：

# use flash attention or a manual implementation?
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
if not self.flash:
    print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
    mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
    mask = torch.triu(mask, diagonal=1)
    self.register_buffer("mask", mask)

Transformer

configure_optimizers 中有一个 use_fused 属性：

use_fused 是一个布尔值，用于决定是否使用融合版本的 AdamW 优化器。融合版本的优化器在某些操作上进行了优化，以提高计算效率和性能。

如何估算大模型的参数

参数介绍：

LLM 模型的参数主要来自于以下几个部分：

参数量 = vocab_size * dim + n_layers * dim * dim * n_heads + dim * num_classes

以 Llama2-Chinese-218M-v1 为例：max_seq_len=1024，dim=1024，n_layers=12，n_heads=8。

参数量 = 64793 * 1024 + 12 * 1024 * 1024 * 8 + 1024 * 100

数据集下载

该模型采用了一些开放数据集，非常有价值，我也一并下载之。

Hugging Face 数据集下载

参见此文：下载Hugging Face的数据集并离线使用的方法 - 知乎：

# 下载
import datasets
dataset = datasets.load_dataset("dataset_name")
dataset.save_to_disk('your_path')

# 加载
import datasets
dataset = load_from_disk("your_path")

在《Downloading datasets》提到，数据集本身就是 Git LFS Repo，直接 Clone 就行，这样更加方便。

我将收集到的数据集记录到。

网络资源

本文作者：Maeiee

本文链接：baby-llama2-chinese 实践笔记

版权声明：如无特别声明，本文即为原创文章，版权归 Maeiee 所有，未经允许不得转载！

喜欢我文章的朋友请随缘打赏，鼓励我创作更多更好的作品！